import os
import cv2
import numpy as np
from rknnlite.api import RKNNLite
import traceback
import sys

# 常量配置
IMG_SIZE = 224
CLASSES = ["moldy", "new"]
RKNN_MODEL = 'best.rknn'
IMG_FOLDER = 'img'

def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()

def preprocess(img):
    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    return np.expand_dims(img, axis=0)

def postprocess(output):
    probs = softmax(output[0]).squeeze()
    class_id = np.argmax(probs).item()
    confidence = probs[class_id].item()
    return class_id, confidence

def rknn_main():
    try:
        print(f"Current working directory: {os.getcwd()}")
        print(f"Python sys.path: {sys.path}")
        
        rknn = RKNNLite()
        results = []

        # 加载RKNN模型
        print('Loading RKNN model...')
        ret = rknn.load_rknn(RKNN_MODEL)
        if ret != 0:
            print('Load RKNN model failed!')
            return ["error"]
        
        # 初始化运行时
        print('Initializing runtime...')
        ret = rknn.init_runtime()
        if ret != 0:
            print('Init runtime failed!')
            return ["error"]
        
        # 获取img文件夹中的第一张图片
        img_files = os.listdir(IMG_FOLDER)
        if not img_files:
            print("Error: No images found in img folder")
            return ["error"]
        
        img_path = os.path.join(IMG_FOLDER, img_files[0])
        print(f"Processing image: {img_path}")
        img = cv2.imread(img_path)
        
        if img is None:
            print(f"Error: Failed to read image {img_path}")
            return ["error"]
        
        # 预处理
        input_data = preprocess(img)
        
        # 执行推理
        outputs = rknn.inference(inputs=[input_data])
        
        # 后处理
        class_id, confidence = postprocess(outputs)
        result = CLASSES[class_id]
        
        print(f"Recognition result: {result}, Confidence: {confidence:.4f}")
        
        # 释放资源
        rknn.release()
        return [result]
    
    except Exception as e:
        print(f"Error in rknn_main: {str(e)}")
        traceback.print_exc()
        return ["error"]

# 测试代码
if __name__ == '__main__':
    print("Running as standalone Python script")
    results = rknn_main()
    print("Final results:", results)



